
// 初始化画板
const canvas = document.getElementById('drawingCanvas');
const ctx = canvas.getContext('2d');
let isDrawing = false;

// 初始化结果图表
const resultCtx = document.getElementById('confidenceChart').getContext('2d');
const resultChart = new Chart(resultCtx, {
  type: 'bar',
  data: {
    labels: ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'],
    datasets: [{
      label: '预测概率',
      data: Array(10).fill(0),
      backgroundColor: 'rgba(124, 58, 237, 0.7)',
      borderColor: 'rgba(124, 58, 237, 1)',
      borderWidth: 1
    }]
  },
  options: {
    responsive: true,
    scales: {
      y: {
        beginAtZero: true,
        max: 1
      }
    }
  }
});

// 画板交互逻辑
function initDrawing() {
  ctx.fillStyle = 'black';
  ctx.fillRect(0, 0, canvas.width, canvas.height);
  ctx.lineWidth = 15;
  ctx.lineCap = 'round';
  ctx.strokeStyle = 'white';
  
  // 触摸设备支持
  canvas.addEventListener('mousedown', startDrawing);
  canvas.addEventListener('mousemove', draw);
  canvas.addEventListener('mouseup', stopDrawing);
  canvas.addEventListener('mouseout', stopDrawing);
  canvas.addEventListener('touchstart', handleTouch);
  canvas.addEventListener('touchmove', handleTouch);
  canvas.addEventListener('touchend', stopDrawing);
}

function handleTouch(e) {
  e.preventDefault();
  const touch = e.touches[0];
  const mouseEvent = new MouseEvent(
    e.type === 'touchstart' ? 'mousedown' : 'mousemove',
    {
      clientX: touch.clientX,
      clientY: touch.clientY
    }
  );
  canvas.dispatchEvent(mouseEvent);
}

function startDrawing(e) {
  isDrawing = true;
  draw(e);
}

function draw(e) {
  if (!isDrawing) return;
  
  const rect = canvas.getBoundingClientRect();
  const x = e.clientX - rect.left;
  const y = e.clientY - rect.top;
  
  ctx.lineTo(x, y);
  ctx.stroke();
  ctx.beginPath();
  ctx.moveTo(x, y);
}

function stopDrawing() {
  isDrawing = false;
  ctx.beginPath();
}

// 清除画板
document.getElementById('clearBtn').addEventListener('click', () => {
  ctx.fillStyle = 'black';
  ctx.fillRect(0, 0, canvas.width, canvas.height);
  resultChart.data.datasets[0].data = Array(10).fill(0);
  resultChart.update();
  document.getElementById('predictionResult').innerHTML = 
    '<p class="text-gray-500">请在手写板绘制数字后点击识别</p>';
});

// 构建CNN模型
function createModel() {
  const model = tf.sequential();
  
  model.add(tf.layers.conv2d({
    inputShape: [28, 28, 1],
    filters: 8,
    kernelSize: 5,
    strides: 1,
    activation: 'relu',
    kernelInitializer: 'varianceScaling'
  }));
  
  model.add(tf.layers.maxPooling2d({
    poolSize: [2, 2],
    strides: [2, 2]
  }));
  
  model.add(tf.layers.conv2d({
    filters: 16,
    kernelSize: 3,
    strides: 1,
    activation: 'relu',
    kernelInitializer: 'varianceScaling'
  }));
  
  model.add(tf.layers.maxPooling2d({
    poolSize: [2, 2],
    strides: [2, 2]
  }));
  
  model.add(tf.layers.flatten());
  
  model.add(tf.layers.dense({
    units: 128,
    activation: 'relu'
  }));
  
  model.add(tf.layers.dense({
    units: 10,
    activation: 'softmax'
  }));
  
  model.compile({
    optimizer: tf.train.adam(0.001),
    loss: 'categoricalCrossentropy',
    metrics: ['accuracy']
  });
  
  return model;
}

let model = createModel();

// 预处理画板图像
function preprocessImage() {
  return tf.tidy(() => {
    const tempCanvas = document.createElement('canvas');
    tempCanvas.width = 28;
    tempCanvas.height = 28;
    const tempCtx = tempCanvas.getContext('2d');
    tempCtx.drawImage(canvas, 0, 0, 28, 28);
    const imgData = tempCtx.getImageData(0, 0, 28, 28);
    
    const tensor = tf.browser.fromPixels(imgData, 1)
      .toFloat()
      .div(255.0)
      .reshape([1, 28, 28, 1]);
    
    return tensor;
  });
}

// 预测数字
document.getElementById('predictBtn').addEventListener('click', async () => {
  const tensor = preprocessImage();
  const predictions = await model.predict(tensor).data();
  
  // 更新图表
  resultChart.data.datasets[0].data = Array.from(predictions);
  resultChart.update();
  
  // 显示最高概率结果
  const maxProb = Math.max(...predictions);
  const predictedNum = predictions.indexOf(maxProb);
  
  document.getElementById('predictionResult').innerHTML = `
    <div class="text-6xl font-bold text-purple-700 mb-2">${predictedNum}</div>
    <div class="text-gray-600">置信度: ${(maxProb * 100).toFixed(1)}%</div>
  `;
  
  tensor.dispose();
});

// 加载MNIST数据
async function loadMNISTData() {
  const data = new MnistData();
  await data.load();
  return data;
}

class MnistData {
  async load() {
    const mnistImages = require('mnist-digits');
    this.trainImages = mnistImages.training.images;
    this.trainLabels = mnistImages.training.labels;
    this.testImages = mnistImages.test.images;
    this.testLabels = mnistImages.test.labels;
  }

  getTrainData() {
    return {
      images: tf.tensor4d(this.trainImages.map(img => 
        img.flatMap(row => row.map(pixel => pixel / 255))
      ), [this.trainImages.length, 28, 28, 1]),
      labels: tf.oneHot(tf.tensor1d(this.trainLabels, 'int32'), 10).toFloat()
    };
  }
}

// 训练模型
document.getElementById('trainBtn').addEventListener('click', async () => {
  const epochs = parseInt(document.getElementById('epochs').value);
  const batchSize = parseInt(document.getElementById('batchSize').value);
  
  const progressContainer = document.getElementById('trainingProgress');
  const progressBar = document.getElementById('progressBar');
  const progressText = document.getElementById('progressText');
  
  progressContainer.classList.remove('hidden');
  
  try {
    const mnistData = await loadMNISTData();
    const {images: trainImages, labels: trainLabels} = mnistData.getTrainData();
    
    await model.fit(trainImages, trainLabels, {
      epochs,
      batchSize,
      validationSplit: 0.1,
      callbacks: {
        onEpochEnd: (epoch, logs) => {
          const progress = ((epoch + 1) / epochs) * 100;
          progressBar.style.width = `${progress}%`;
          progressText.textContent = `${Math.round(progress)}%`;
          console.log(`Epoch ${epoch + 1}: loss = ${logs.loss.toFixed(4)}`);
        }
      }
    });
    
    alert('模型训练完成！');
  } catch (error) {
    console.error('训练出错:', error);
    alert('训练出错: ' + error.message);
  } finally {
    progressBar.style.width = '0%';
    progressText.textContent = '0%';
  }
});

// 初始化应用
initDrawing();
